import os
import re
import PyPDF2
from pathlib import Path

class PDFMergerTool:
    @staticmethod
    def merge_pdfs(pdf_list, output_filename):
        """
        合并多个 PDF 文件
        :param pdf_list: PDF 文件路径列表
        :param output_filename: 输出的合并 PDF 文件名
        """
        pdf_merger = PyPDF2.PdfMerger()

        for pdf in pdf_list:
            pdf_path = Path(pdf)  # 确保路径格式正确
            if not pdf_path.exists():
                print(f"错误: 找不到文件 {pdf_path}")
                continue  # 跳过不存在的文件

            print(f"正在合并文件: {pdf_path}")
            pdf_merger.append(str(pdf_path))  # 转换为字符串路径

        output_path = Path(output_filename)
        with output_path.open('wb') as output_pdf:
            pdf_merger.write(output_pdf)
        pdf_merger.close()  # 释放资源

        print(f"合并完成，输出文件: {output_path}")

    @staticmethod
    def extract_number(filename):
        """提取文件名前的数字部分"""
        match = re.match(r'(\d+)', Path(filename).stem)
        return int(match.group(1)) if match else float('inf')

    @staticmethod
    def clean_filename(filename, max_length=255):
        """去除文件名中的非法字符，并保证文件名合理"""
        if isinstance(filename, Path):  # 确保 filename 是字符串
            filename = str(filename)
        # 替换非法字符，合并多个非法字符为单个 '_'
        cleaned = re.sub(r'[\\/:*?"<>|]+', '_', filename)
        # 去除首尾 '_'
        cleaned = cleaned.strip('_')
        # 限制文件名长度
        return cleaned[:max_length]


    @classmethod
    def split_and_merge_pdfs(cls, folder_path, prefix="merged_output", delete_after_merge=False):
        """
        读取文件夹中的所有 PDF 文件，并根据数量进行合并
        仅合并文件名为纯数字的 PDF 文件。
        """
        prefix = cls.clean_filename(prefix)
        folder = Path(folder_path)  # 统一路径格式
        if not folder.exists():
            print(f"错误: 目录 {folder} 不存在")
            return

        # 仅筛选数字命名的 PDF 文件
        pdf_files = [f for f in folder.iterdir() if f.suffix == '.pdf' and re.fullmatch(r'\d+\.pdf', f.name)]
        pdf_files.sort(key=lambda f: cls.extract_number(f.name))
        num_files = len(pdf_files)

        print(f"找到 {num_files} 个符合条件的 PDF 文件。")

        if num_files == 0:
            print("没有找到符合条件的 PDF 文件，操作终止。")
            return

        merged_files = []

        if num_files <= 50:
            print("文件数量不超过50，直接合并所有文件。")
            merged_filename = str(Path(folder / f"{prefix}.pdf"))
            merged_files.append(merged_filename)
            cls.merge_pdfs([str(f) for f in pdf_files], str(merged_filename))

        elif num_files <= 110:
            part_size = num_files // 2
            print("文件数量超过50但不超过110，分成上下两部分进行合并。")
            merged_files.append(str(Path(folder / f"{prefix}_part_1.pdf")))
            merged_files.append(str(Path(folder / f"{prefix}_part_2.pdf")))

            cls.merge_pdfs([str(f) for f in pdf_files[:part_size]], str(merged_files[-2]))
            cls.merge_pdfs([str(f) for f in pdf_files[part_size:]], str(merged_files[-1]))

        else:
            part_size = num_files // 3
            remainder = num_files % 3
            print("文件数量超过110，分成上、中、下三部分进行合并。")

            first_part = pdf_files[:part_size + (1 if remainder > 0 else 0)]
            second_part = pdf_files[part_size + (1 if remainder > 0 else 0): 2 * part_size + (1 if remainder > 1 else 0)]
            third_part = pdf_files[2 * part_size + (1 if remainder > 1 else 0):]

            merged_files.append(str(Path(folder / f"{prefix}_part_1.pdf")))
            merged_files.append(str(Path(folder / f"{prefix}_part_2.pdf")))
            merged_files.append(str(Path(folder / f"{prefix}_part_3.pdf")))

            cls.merge_pdfs([str(f) for f in first_part], str(merged_files[-3]))
            cls.merge_pdfs([str(f) for f in second_part], str(merged_files[-2]))
            cls.merge_pdfs([str(f) for f in third_part], str(merged_files[-1]))

        if delete_after_merge:
            import time
            time.sleep(2)  # 等待文件释放
            for pdf in pdf_files:
                try:
                    pdf.unlink()
                    print(f"已删除: {pdf}")
                except PermissionError as e:
                    print(f"无法删除 {pdf}: {e}")
            print("删除完成")

